package LibDL.models;

import LibDL.core.*;
import LibDL.core.nn.*;
import LibDL.core.nn.Module;

import java.util.ArrayList;
import java.util.List;

public class NCF extends Module {
    Embedding embed_user_GMF, embed_item_GMF, embed_user_MLP, embed_item_MLP;
    Linear predict_layer;
    ArrayList<Dropout> dropout_list = new ArrayList<>();
    ArrayList<Linear> linear_list = new ArrayList<>();

    public NCF(int num_user, int num_item, int num_factor, int num_layers, double dropout_rate) {
        this.embed_user_GMF = new Embedding(num_user, num_factor);
        this.embed_item_GMF = new Embedding(num_item, num_factor);
        this.embed_user_MLP = new Embedding(num_user, num_factor * (2 << (num_layers - 1)));
        this.embed_item_MLP = new Embedding(num_item, num_factor * (2 << (num_layers - 1)));

        for (int i = 0; i < num_layers; i++) {
            int input_size = num_factor * 2 * (2 << (num_layers - i));
            this.dropout_list.add(new Dropout(dropout_rate));
            this.linear_list.add(new Linear(input_size, input_size / 2));
        }

        this.predict_layer = new Linear(num_factor, 1);

        init_weight();
    }

    public void init_weight() {
        try (NoGrad noGrad = new NoGrad()) {
            init.normal_(this.embed_user_GMF.getWeight(), 0, 0.01);
            init.normal_(this.embed_user_MLP.getWeight(), 0, 0.01);
            init.normal_(this.embed_item_GMF.getWeight(), 0, 0.01);
            init.normal_(this.embed_item_MLP.getWeight(), 0, 0.01);

            for (Linear linear : linear_list) {
                init.xavier_uniform_(linear.getWeight());
            }
            init.kaiming_uniform_(this.predict_layer.getWeight(), 1, init.FanMode.FanIn, init.Nonlinearity.Sigmoid);
        }
    }

    public Tensor list(Tensor input) {
        for (int i = 0; i < linear_list.size(); i++) {
            input = dropout_list.get(i).forward(input);
            input = linear_list.get(i).forward(input);
            input = functional.relu(input);
        }
        return input;
    }

    public Tensor forward(Tensor user, Tensor item) {
        Tensor res_user_GMF = embed_user_GMF.forward(user);
        Tensor res_item_GMF = embed_item_GMF.forward(item);
        Tensor output_GMF = res_user_GMF.mul(res_item_GMF);

        Tensor res_user_MLP = embed_user_MLP.forward(user);
        Tensor res_item_MLP = embed_item_MLP.forward(item);
        Tensor interaction = functional.cat(new StdVectorTensor(new Tensor[]{res_user_MLP, res_item_MLP}), 1);
        Tensor output_MLP = list(interaction);

        return predict_layer.forward(
                functional.cat(new StdVectorTensor(new Tensor[]{output_GMF, output_MLP}), 1)
        );
    }
}